#include "sanitizer_printf.h"
#include <stdarg.h>
#include "asan_interceptors.h"
#include <unistd.h>
#include "common_flags.h"
#include "sanitizer_libc.h"
#include "asan_report.h"

namespace __sanitizer
{

const int SAFE_PRINTF_BUFF_SIZE = 4096 * 4;

bool ColorizeReports() {
  const char *flag = common_flags()->color;
  return internal_strcmp(flag, "always") == 0 || \
         internal_strcmp(flag, "auto") == 0;
}

const char* PrintfDecorator::GetShadowByteStr(su_t v) {
    #define COLOR_SD_SNPRINTF(color) \
      REAL(snprintf(shadow_byte_buff, 31, "%s%x%s", color(), v & 0x0ff, Default()));
    switch (v) {
    case kAsanHeapLeftRedzoneMagic:
      COLOR_SD_SNPRINTF(Red); break;
    case kAsanHeapFreeMagic:
      COLOR_SD_SNPRINTF(Magenta); break;
    case kAsanStackLeftRedzoneMagic:
      COLOR_SD_SNPRINTF(Red); break;
    case kAsanStackMidRedzoneMagic:
      COLOR_SD_SNPRINTF(Red); break;
    case kAsanStackRightRedzoneMagic:
      COLOR_SD_SNPRINTF(Red); break;
    case kAsanStackAfterReturnMagic:
      COLOR_SD_SNPRINTF(Magenta); break;
    case kAsanInitializationOrderMagic:
      COLOR_SD_SNPRINTF(Cyan); break;
    case kAsanUserPoisonedMemoryMagic:
      COLOR_SD_SNPRINTF(Blue); break;
    case kAsanContiguousContainerOOBMagic:
      COLOR_SD_SNPRINTF(Blue); break;
    case kAsanStackUseAfterScopeMagic:
      COLOR_SD_SNPRINTF(Magenta); break;
    case kAsanInternalHeapMagic:
      COLOR_SD_SNPRINTF(Yellow); break;
    case kAsanGlobalRedzoneMagic:
      COLOR_SD_SNPRINTF(Red); break;
    case kAsanArrayCookieMagic:
      COLOR_SD_SNPRINTF(Red); break;
    case kAsanIntraObjectRedzone:
      COLOR_SD_SNPRINTF(Yellow); break;
    case kAsanAllocaLeftMagic:
      COLOR_SD_SNPRINTF(Blue); break;
    case kAsanAllocaRightMagic:
      COLOR_SD_SNPRINTF(Blue); break;
    default:
      if (0 <= v && v <= 0x0f) {
        REAL(snprintf(shadow_byte_buff, 31, "0%x", v & 0x0ff));
      }
      else {
        REAL(snprintf(shadow_byte_buff, 31, "%x", v & 0x0ff));
      }
      break;
    }
    #undef COLOR_SD_SNPRINTF
    return shadow_byte_buff;
}

int safe_printf(const char* format, ...) {
  va_list args;
  va_start(args, format);
  char * buff = (char*)REAL(malloc(SAFE_PRINTF_BUFF_SIZE + 1));
  int char_count_or_error = vsnprintf(buff, SAFE_PRINTF_BUFF_SIZE, format, args);
  va_end(args);

  if (char_count_or_error < 0 || char_count_or_error > SAFE_PRINTF_BUFF_SIZE) {
    REAL(free(buff));
    return char_count_or_error;
  }
  buff[SAFE_PRINTF_BUFF_SIZE] = '\0';
  write(2, buff, char_count_or_error);
  REAL(free(buff));
  return char_count_or_error;
}

void Printf(const char *format, ...) {
  va_list args;
  va_start(args, format);
  // The following is the code of safe_printf
  // Call safe_printf here result in errornous output
  char * buff = (char*)REAL(malloc(SAFE_PRINTF_BUFF_SIZE + 1));
  int char_count_or_error = vsnprintf(buff, SAFE_PRINTF_BUFF_SIZE, format, args);

  if (char_count_or_error < 0 || char_count_or_error > SAFE_PRINTF_BUFF_SIZE) {
    REAL(free(buff));
    va_end(args);
    return;
  }
  buff[SAFE_PRINTF_BUFF_SIZE] = '\0';
  write(2, buff, char_count_or_error);
  REAL(free(buff));
  va_end(args);
}

void RawWrite(const char* msg) {
  uptr read_len = internal_strnlen(msg, 4096);
  write(2, msg, read_len);
}

void CheckPrintfVars(const char* format, va_list args) {
  // Go through char* pointers of the format string
  static const char *kPrintfFormatsHelp =
      "Supported Printf formats: \%([0-9]*)?(z|l|ll|L)?{d,u,x,X,e,E,g,G,a,A,f,F}; \%p; "
      "\%[-]([0-9]*)?(\\.\\*)?s; \%c\nProvided format: ";
  const char* cur = format;
  char pre = 0;

  int result = 0;
  for (; *cur; cur++) {
    if (*cur != '%') {
      continue;
    }
    cur++;
    bool left_justified = *cur == '-';
    if (left_justified)
      cur++;
    bool have_width = (*cur >= '0' && *cur <= '9');
    bool pad_with_zero = (*cur == '0');
    int width = 0;
    if (have_width) {
      while (*cur >= '0' && *cur <= '9') {
        width = width * 10 + *cur++ - '0';
      }
    }
    bool have_precision = (cur[0] == '.' && cur[1] == '*');
    int precision = -1;
    if (have_precision) {
      cur += 2;
      precision = va_arg(args, int);
    }
    bool have_z = (*cur == 'z');
    cur += have_z;
    bool have_l = (cur[0] == 'l' && cur[1] != 'l');
    cur += have_l;
    bool have_L = (cur[0] == 'L');
    cur += have_L;
    bool have_ll = (cur[0] == 'l' && cur[1] == 'l');
    cur += have_ll * 2;
    // %Lf is %llf
    if (have_L) have_ll = true;
    const bool have_length = have_z || have_l || have_ll;
    const bool have_flags = have_width || have_length;
    // At the moment only %s supports precision and left-justification.
    CHECK(!((precision >= 0 || left_justified) && *cur != 's'));
    switch (*cur) {
      case 'd': {
        s64 dval = have_ll  ? va_arg(args, s64)
                   : have_z ? va_arg(args, sptr)
                   : have_l ? va_arg(args, long)
                            : va_arg(args, int);
        break;
      }
      case 'u':
      case 'x':
      case 'X': {
        u64 uval = have_ll  ? va_arg(args, u64)
                   : have_z ? va_arg(args, uptr)
                   : have_l ? va_arg(args, unsigned long)
                            : va_arg(args, unsigned);
        bool uppercase = (*cur == 'X');
        break;
      }
      case 'e' || 'E' || 'g' || 'G' || 'a' || 'A' || 'f' || 'F': {
        // It is said that C99 will auto upscale float to double
        long double uval = have_ll ? va_arg(args, long double)
                                  : have_l ? va_arg(args, double)
                                           : va_arg(args, double);
        break;
      }
      case 'p': {
        // RAW_CHECK_VA(!have_flags, kPrintfFormatsHelp, format);
        va_arg(args, uptr);
        break;
      }
      case 's': {
        // RAW_CHECK_VA(!have_length, kPrintfFormatsHelp, format);
        // Only left-justified width is supported.
        CHECK(!have_width || left_justified);
        char* str = va_arg(args, char*);
        CheckAndReport((void*)str, 1, false);
        uptr tmp_read_len = REAL(strlen(str));
        CheckAndReport((void*)str, (tmp_read_len + 1), false);
        break;
      }
      case 'c': {
        va_arg(args, int);
        break;
      }
      case '%' : {
        // RAW_CHECK_VA(!have_flags, kPrintfFormatsHelp, format);
        break;
      }
      default: {
        // RAW_CHECK_VA(false, kPrintfFormatsHelp, format);
      }
    }
  }
}

} // namespace __sanitizer

